package com.sxd.bill.ner.utils;

import java.util.ArrayList;
import java.util.List;

import com.sxd.bill.common.enums.LabelNameEnum;
import com.sxd.bill.common.item.LabelItem;
import com.sxd.bill.common.item.TextLabelItem;

import lombok.Data;

public class OutputUtility {

    public static List<LabelChunk> getLabelChunks(List<String> seq) {
        List<LabelChunk> chunks = new ArrayList<>();
        LabelChunk chunk = new LabelChunk();
        for (int i = 0; i < seq.size(); i++) {
            String tag = seq.get(i);
            if (tag.startsWith("S-")) {
                if (chunk.getEndIndex() != -1) {
                    chunks.add(chunk);
                }
                String labelName = tag.replaceFirst("^S\\-", "").trim();
                chunks.add(new LabelChunk(labelName, i, i));
                chunk = new LabelChunk();
                continue;
            }
            if (tag.startsWith("B-")) {
                if (chunk.getEndIndex() != -1) {
                    chunks.add(chunk);
                }
                String labelName = tag.replaceFirst("^B\\-", "").trim();
                chunk = new LabelChunk(labelName, i, -1);
            } else if (tag.startsWith("I-") && chunk.getStartIndex() != -1) {
                String tmpType = tag.replaceFirst("I\\-", "").trim();
                if (tmpType.equals(chunk.getLabelName())) {
                    chunk.setEndIndex(i);
                }
                if (i == seq.size() - 1) {
                    chunks.add(chunk);
                }
            } else {
                if (chunk.getEndIndex() != -1) {
                    chunks.add(chunk);
                }
                chunk = new LabelChunk();
            }
        }
        return chunks;
    }

    public static TextLabelItem getTextLabelItem(List<String> tagSequences, String Text) {
        List<LabelChunk> chunks = getLabelChunks(tagSequences);
        List<LabelItem> labels = new ArrayList<>();
        for (LabelChunk chunk : chunks) {
            String labelNameString = chunk.getLabelName();
            LabelNameEnum labelNameEnum;
            try {
                labelNameEnum = LabelNameEnum.fromCnName(labelNameString);
            } catch (Exception e) {
                // e.printStackTrace();
                System.out.println(e.getMessage()); // 或者记录日志
                continue;
            }
            LabelItem labelItem = new LabelItem();
            labelItem.setLabelName(labelNameEnum);

            int startIndex = chunk.getStartIndex();
            int endIndex = chunk.getEndIndex() + 1; // 索引选取时左闭右开，故+1

            labelItem.setStartIndex(startIndex);
            labelItem.setEndIndex(endIndex);

            String entity = Text.substring(startIndex, endIndex);
            labelItem.setLabelEntity(entity);

            labels.add(labelItem);
        }
        return new TextLabelItem(Text, labels);
    }
}

@Data
class LabelChunk {
    private String labelName;
    private int startIndex;
    private int endIndex;

    public LabelChunk() {
        this("UNK", -1, -1);
    }

    public LabelChunk(String labelName, int startIndex, int endIndex) {
        this.labelName = labelName;
        this.startIndex = startIndex;
        this.endIndex = endIndex;
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj) {
            return true;
        }
        if (obj == null || getClass() != obj.getClass()) {
            return false;
        }
        LabelChunk chunk = (LabelChunk) obj;
        return labelName.equals(chunk.labelName) && startIndex == chunk.startIndex && endIndex == chunk.endIndex;
    }

    @Override
    public String toString() {
        return String.format("LabelChunk(labelName=%s,startIndex=%d, endIndex=%d)", labelName, startIndex, endIndex);
    }
}